# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     datalist
   Description :   
   Author :       lth
   date：          2022/12/9
-------------------------------------------------
   Change Activity:
                   2022/12/9 9:16: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import os

import torch
import torchvision


def rescale(x, lo, hi):
    """Rescale a tensor to [lo,hi]."""
    assert (lo < hi), "[rescale] lo={0} must be smaller than hi={1}".format(lo, hi)
    old_width = torch.max(x) - torch.min(x)
    old_center = torch.min(x) + (old_width / 2.)
    new_width = float(hi - lo)
    new_center = lo + (new_width / 2.)
    # shift everything back to zero:
    x = x - old_center
    # rescale to correct width:
    x = x * (new_width / old_width)
    # shift everything to the new center:
    x = x + new_center
    # return:
    return x


def load_cifar10(train=True):
    """Rescale and preprocess CIFAR10 dataset."""
    # check if ZCA matrix exists on dataset yet:
    assert os.path.exists("./dataset/cifar/zca_matrix.pt"), \
        "[load_cifar10] ZCA whitening matrix not built! Run `python make_datasets.py` first."
    zca_matrix = torch.load("./dataset/cifar/zca_matrix.pt")

    cifar10_transform = torchvision.transforms.Compose([
        # convert PIL image to tensor:
        torchvision.transforms.ToTensor(),
        # flatten:
        # torchvision.transforms.Lambda(lambda x: x.view(-1)),
        # add uniform noise ~ [-1/256, +1/256]:
        # torchvision.transforms.Lambda(lambda x: (x + torch.rand_like(x).div_(128.).add_(-1. / 256.))),
        # rescale to [-1,1]:
        # torchvision.transforms.Lambda(lambda x: rescale(x, -1., 1.)),
        # exact ZCA:
        # torchvision.transforms.LinearTransformation(zca_matrix,torch.zeros_like(zca_matrix))
    ])
    return torchvision.datasets.CIFAR10(root="./dataset/cifar", train=train, transform=cifar10_transform,
                                        download=False),zca_matrix

from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

train_trasform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

test_trasform = transforms.Compose([
    transforms.ToTensor(),
    # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])


class CelebAHQ(Dataset):
    def __init__(self, path="E:\Datasets2\CelebAMask-HQ\CelebA-HQ-img", mode="train"):
        super(CelebAHQ, self).__init__()
        self.data = self.get_data(path)
        self.image_size = (64,64)
        self.mode = mode
    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        path = self.data[index]
        img = Image.open(path).convert("RGB")
        img = img.resize(self.image_size)

        if self.mode == "train":
            return train_trasform(img)
        else:
            return test_trasform(img)

    @staticmethod
    def get_data(path):
        data = []
        for root, dir, files in os.walk(path):
            for file in files:
                data.append(os.path.join(root, file))
        return data



if __name__ == "__main__":
    load_cifar10()
